﻿//#include <Windows.h>
//#include <winternl.h>
#include <fast_io.h>
#include "../thread_freezer.hpp"

namespace safetyhook {
void execute_while_frozen(
	const std::function<void()>& run_fn, const std::function<void(uint32_t, void*, fast_io::win32::context&)>& visit_fn) {
    // Freeze all threads.
    int num_threads_frozen;
    auto first_run = true;

    do {
        num_threads_frozen = 0;
        void* thread{};

        while (true) {
            void* next_thread{};
			const auto status = fast_io::win32::nt::NtGetNextThread(fast_io::win32::GetCurrentProcess(), thread, 0x800 | 0x2 | 0x8 | 0x10, 0, 0, &next_thread);

            if (thread != nullptr) {
                fast_io::win32::CloseHandle(thread);
            }

            if (status < 0) {
                break;
            }

            thread = next_thread;

            const auto thread_id = fast_io::win32::GetThreadId(thread);

            if (thread_id == 0 || thread_id == fast_io::win32::GetCurrentThreadId()) {
                continue;
            }

            const auto suspend_count = fast_io::win32::SuspendThread(thread);

            if (suspend_count == static_cast<::std::uint_least32_t>(-1)) {
                continue;
            }

            // Check if the thread was already frozen. Only resume if the thread was already frozen, and it wasn't the
            // first run of this freeze loop to account for threads that may have already been frozen for other reasons.
            if (suspend_count != 0 && !first_run) {
				fast_io::win32::ResumeThread(thread);
                continue;
            }

            fast_io::win32::context thread_ctx{};

            thread_ctx.ContextFlags = (0x00100000L | 0x00000001L) | (0x00100000L | 0x00000002L) | (0x00100000L | 0x00000008L);

            if (!fast_io::win32::GetThreadContext(thread, &thread_ctx)) {
                continue;
            }

            if (visit_fn) {
                visit_fn(thread_id, thread, thread_ctx);
            }

            ++num_threads_frozen;
        }

        first_run = false;
    } while (num_threads_frozen != 0);

    // Run the function.
    if (run_fn) {
        run_fn();
    }

    // Resume all threads.
    void* thread{};

    while (true) {
        void* next_thread{};
       

        const auto status = fast_io::win32::nt::NtGetNextThread(fast_io::win32::GetCurrentProcess(), thread, 0x800 | 0x2 | 0x8 | 0x10, 0, 0, &next_thread);

        if (thread != nullptr) {
            fast_io::win32::CloseHandle(thread);
        }

        if (status < 0) {
            break;
        }

        thread = next_thread;

        const auto thread_id = fast_io::win32::GetThreadId(thread);

        if (thread_id == 0 || thread_id == fast_io::win32::GetCurrentThreadId()) {
            continue;
        }

        fast_io::win32::ResumeThread(thread);
    }
}

void fix_ip(fast_io::win32::context& ctx, uint8_t* old_ip, uint8_t* new_ip) {
#ifdef _M_X64
    auto ip = ctx.Rip;
#else
    auto ip = ctx.Eip;
#endif

    if (ip == reinterpret_cast<uintptr_t>(old_ip)) {
        ip = reinterpret_cast<uintptr_t>(new_ip);
    }

#ifdef _M_X64
    ctx.Rip = ip;
#else
    ctx.Eip = ip;
#endif
}
} // namespace safetyhook